from tqdm import tqdm
import json
import pickle
from scripts.server.chat_api import chat
from tqdm import tqdm
import re

with open('./entities_new.pkl','rb') as f:
    entities = pickle.load(f)
with open('./relationships_new.pkl','rb') as f:
    relationships = pickle.load(f)
with open('./record.json','r') as f:
    record = json.load(f)


COMMUNITY_REPORT_PROMPT = """
您是一位人工智能助手，帮助人类分析师进行一般信息发现。信息发现是识别和评估与网络中某些实体（例如组织和个人）相关的信息的过程。

# 目标
根据属于社区的实体列表以及它们之间的关系和可选的相关声明，编写一份全面的社区报告。报告将用于向决策者提供与社区及其潜在影响相关的信息。报告内容包括社区关键实体的概述、它们的合法合规性、技术能力、声誉和值得注意的声明。

# 报告结构
报告应包括以下部分：
标题：代表其关键实体的社区名称 - 标题应简短但具体。如果可能，包括代表性的命名实体在标题中。
摘要：社区整体结构的执行摘要，其实体如何相互关联，以及与其实体相关的重大信息。
影响严重性评级：一个0-10之间的浮点分数，代表社区内实体构成的影响严重性。影响是社区重要性的评分。
评级解释：给出影响严重性评级的一句话解释。
详细发现：关于社区的5-10个关键见解的列表。每个见解都应有一个简短的摘要，然后是几段根据下面的依据规则的解释性文本。要全面。

以格式良好的JSON格式字符串返回输出，格式如下：
{{
"title": <报告标题>,
"summary": <执行摘要>,
"rating": <影响严重性评级>,
"rating_explanation": <评级解释>,
"findings": [
{{
"summary":<见解1摘要>,
"explanation": <见解1解释>
}},
{{
"summary":<见解2摘要>,
"explanation": <见解2解释>
}}
]
}}


# 示例输入
文本：

实体
id,实体,描述
5,VERDANT OASIS PLAZA,Verdant Oasis Plaza是Unity March的地点
6,HARMONY ASSEMBLY,Harmony Assembly是一个在Verdant Oasis Plaza举行游行的组织

关系
id,源,目标,描述
37,VERDANT OASIS PLAZA,UNITY MARCH,Verdant Oasis Plaza是Unity March的地点
38,VERDANT OASIS PLAZA,HARMONY ASSEMBLY,Harmony Assembly在Verdant Oasis Plaza举行游行
39,VERDANT OASIS PLAZA,UNITY MARCH,Unity March在Verdant Oasis Plaza举行
40,VERDANT OASIS PLAZA,TRIBUNE SPOTLIGHT,Tribune Spotlight正在报道在Verdant Oasis Plaza举行的Unity游行
41,VERDANT OASIS PLAZA,BAILEY ASADI,Bailey Asadi在Verdant Oasis Plaza就游行发表演讲
43,HARMONY ASSEMBLY,UNITY MARCH,Harmony Assembly正在组织Unity March

输出：
{{
"title": "Verdant Oasis Plaza和Unity March",
"summary": "社区围绕Verdant Oasis Plaza展开，该广场是Unity March的地点。该广场与Harmony Assembly、Unity March和Tribune Spotlight有关联，所有这些都与游行事件有关。",
"rating": 5.0,
"rating_explanation": "由于Unity March期间可能发生动荡或冲突，因此影响严重性评级为中等。",
"findings": [
{{
"summary": "Verdant Oasis Plaza作为中心地点",
"explanation": "Verdant Oasis Plaza是这个社区的中心实体，作为Unity March的地点。这个广场是所有其他实体的共同联系，表明它在社区中的重要性。根据游行的性质和它引发反应，该广场与游行的关联可能会导致公共混乱或冲突等问题。"
}},
{{
"summary": "Harmony Assembly在社区中的角色",
"explanation": "Harmony Assembly是这个社区的另一个关键实体，是在Verdant Oasis Plaza组织游行的组织。Harmony Assembly的性质和它的游行可能是潜在威胁的来源，这取决于他们的目标和他们引发反应。Harmony Assembly与广场的关系对于理解这个社区的动态至关重要。"
}},
{{
"summary": "Unity March作为重大事件",
"explanation": "Unity March是在Verdant Oasis Plaza举行的重大事件。这个事件是社区动态的关键因素，根据游行的性质和它引发反应，可能是潜在威胁的来源。游行与广场的关系对于理解这个社区的动态至关重要。"
}},
{{
"summary": "Tribune Spotlight的角色",
"explanation": "Tribune Spotlight正在报道在Verdant Oasis Plaza举行的Unity March。这表明该事件已经吸引了媒体的注意，这可能会扩大其对社区的影响。Tribune Spotlight的角色在塑造公众对事件和涉及实体的看法方面可能非常重要。"
}}
]
}}


# 实际数据
使用以下文本回答你的问题。不要在你的回答中编造任何内容。

文本：
{input_text}

报告应包括以下部分：

标题：代表其关键实体的社区名称 - 标题应简短但具体。尽可能在标题中包含代表性的命名实体。
摘要：社区整体结构的执行摘要，其实体之间的关系，以及与其实体相关的重要信息。
影响严重程度评分：一个介于0-10之间的浮点分数，代表社区内实体构成的影响严重程度。影响是社区得分的重要性。
评分解释：给出影响严重程度评分的单句解释。
详细发现：列出5-10个关于社区的关键见解。每个见解都应有一个简短的摘要，然后是多段基于以下地面规则的解释性文本。要全面。
以以下格式返回格式良好的JSON格式字符串：
{{
"title": <report_title>,
"summary": <executive_summary>,
"rating": <impact_severity_rating>,
"rating_explanation": <rating_explanation>,
"findings": [
{{
"summary":<insight_1_summary>,
"explanation": <insight_1_explanation>
}},
{{
"summary":<insight_2_summary>,
"explanation": <insight_2_explanation>
}}
]
}}

输出：
"""

# class Community():
#     def __init__(self, nodes, label, report):
#         self.nodes = nodes
#         self.label = label
#         self.report = report


if __name__ == '__main__':
    # 获得多层社区关系 `total_community`
    total_community = []
    for floor in range(len(list(record.values())[0])):      # 遍历三层
        labels = []
        for entity_name, group in record.items():
            # 将第 floor 层的归为一起
            label = '-'.join([str(tmp) for tmp in group[:floor+1]])
            labels.append(label)
        labels = list(set(labels))
        total_community.append(labels)

    # 基于社区获取对应实体、边关系
    total_community_detail = []
    for level_community in total_community:
        for community in level_community:
            labels = [int(tmp) for tmp in community.split('-')]     # 社区标签
            # 提取该社区标签的所有实体和边
            community_tmp,  entities_tmp, relations_tmp = None, [], []
            for entity, entity_community in record.items():
                if entity_community[:len(labels)] == labels:
                    community_tmp = labels  # 社区赋值
                    # 找到实体
                    for entity_search in entities:
                        if entity_search[0] == entity:
                            entities_tmp.append(entity_search)
                    # 找到关系
                    for relationship_search in relationships:
                        if relationship_search[0] == entity or relationship_search[1] == entity:
                            relations_tmp.append(relationship_search)
            total_community_detail.append(
                {"community": community_tmp, "entities": entities_tmp, "relations": relations_tmp})


    # 获取社区报告
    for i, community_detail in enumerate(tqdm(total_community_detail)):
        tmp_entities = community_detail['entities']
        tmp_relations = community_detail['relations']

        tmp_entity_input = '\n'.join([f"{i},{tmp_entities[i][0]},{tmp_entities[i][-1]}" for i in range(len(tmp_entities))])
        tmp_relation_input = '\n'.join([f"{i},{tmp_relations[i][0]},{tmp_relations[i][1]},{tmp_relations[i][2]}" for i in range(len(tmp_relations))])

        input_text = f'''实体\nid,实体,描述\n{tmp_entity_input}\n\n\n关系\nid,源,目标,描述\n{tmp_relation_input}'''

        model_input = COMMUNITY_REPORT_PROMPT.replace('{input_text}', input_text)
        answer = chat(model_input,[],0)

        total_community_detail[i]['report'] = answer

    # 保存记录
    with open('total_community_detail.json','w',encoding='utf-8') as f:
        json.dump(total_community_detail, f, ensure_ascii=False, indent=4)